import torch
import os
import torch.nn as nn  # 这里将重要的库命名。
# import torchvision.transforms as transforms
# from torch.utils.data import DataLoader  # 这里引入一个dataloader.
# from torch.utils.data import Dataset  # 这里导入Data
# from torchvision import datasets
# from torch.utils.tensorboard import SummaryWriter
# import torchvision.models as models
import class_function as cf  # 自定义的类和函数。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 初始化GPU,将数据放在GPU上面进行训练。
main_file = os.listdir("./dataset/train/")

print("正在进行的是数据增强")
for j in main_file:
    names = "./dataset/train/" + j
    dataset = cf.MyData(names)
    len_data = dataset.__len__()  # 这里是提取出对应文件夹中图片数量。
    # print(len_data)
    if len_data > 400:
        continue
    else:
        # print(dataset)
        for i in range(int(len_data * 0.5)):  # 每个分类都增加150个数据样本。
            filenames = dataset.get_name()
            # print(filenames)
            image = dataset.__getitem__(i)
            new_image1 = cf.v_function(image)
            new_image1 = new_image1.convert('RGB')
            new_image1.save(filenames + "/new1" + str(i) + ".jpg")

            new_image2 = cf.h_function(image)
            new_image2 = new_image2.convert('RGB')
            new_image2.save(filenames + "/new2" + str(i) + ".jpg")
print("数据增强完成。")

# 接下来是获取相应的数据集。
train, test = cf.tr_te()

# 初始化一些模型和参数。。
num = len(main_file)  # 这里的num指的是分类个数。
net = cf.Re_net(num)  # 进行实例化
net = net.to(device)  # 有GPU就弄到GPU上面去。
loss_function = nn.CrossEntropyLoss().to(device)  # 交叉熵损失函数。
optimizer = torch.optim.Adam(net.parameters(), lr=0.0001)  # 定义优化方法。

# 开始训练。
epoch = 10  # 迭代epoch次即可。
net, train_list, test_list, loss_list = cf.train_function(net, epoch, optimizer, loss_function, train, test)
cf.test_function(net, test)

# 保存模型。
torch.save(net, './net1.pt')
torch.save(net.state_dict(), './W_net1.pt')
